from ultralytics.yolo.utils import ops

import numpy as np
import torch
import cv2


"""
说明：通过深度学习-目标检测 检测孔洞+数字
"""

hole_model = None
weights_hole = "weights/holev1.pt"
stride = 32
imgsz = 640
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dohalf = False
conf_thres = 0.65
iou_thres = 0.6
names = None
color = [255, 255, 0]


def sort_by_confidence(box_list):
    sorted_lst = sorted(box_list, key=lambda x: x[4], reverse=True)
    tuple_lst = [list(item) for item in sorted_lst]
    return tuple_lst


def calculate_intersection_over_union(box1, box2):
    # 计算交集区域的左上角和右下角坐标
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    if x1 < x2 and y1 < y2:
        # 计算交集区域的宽度和高度
        intersection_width = max(0, x2 - x1)
        intersection_height = max(0, y2 - y1)
        # 计算交集区域的面积
        intersection_area = intersection_width * intersection_height
        # 计算两个目标框的面积
        area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
        area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
        # 计算并集区域的面积
        union_area = area1 + area2 - intersection_area
        # 计算 IoU
        iou = intersection_area / union_area
        return iou
    else:
        return 0


def det_hole(img0: np.ndarray, model_obj, labels_, draw: bool = False):
    """
    获取孔洞识别模型结果
    :param img0: 图片的numpy数据格式
    :param weight: 权重文件路径
    :param draw: 是否标图
    :return: 返回预测结果
    """
    nums = []
    with torch.inference_mode():
        img0, img = preprocess(img0)
        hole_model = model_obj
        hole_names = labels_
        preds = hole_model.model(img, augment=False, visualize=False)
        results = postprocess(preds, img, img0)
        # print("results： ", results)
    if results is not None:
        for *xyxy, conf, cls in results:
            nums.append((hole_names[int(cls)], conf, xyxy))
    holes = trim_holes(nums)
    # print("孔洞识别：", holes)
    if draw:
        for _, _, xyxy in holes:
            c1, c2 = (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3]))
            cv2.rectangle(img0, c1, c2, color, thickness=1, lineType=cv2.LINE_AA)
    return holes


def preprocess(img0: np.ndarray, auto=True):
    """
    图像转换
    :param img0:
    :param auto:
    :return:
    """
    img = letterbox(img0, imgsz, stride=stride, auto=auto)[0]

    # Convert
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device, non_blocking=True)
    if dohalf:
        img = img.half()
    else:
        img = img.float()
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)
    # 原始图片，处理后图片
    return img0, img


def postprocess(preds, img, orig_img):
    preds = ops.non_max_suppression(preds, conf_thres, iou_thres, agnostic=False, max_det=300, classes=None, multi_label = True)
    if len(preds) == 0:
        return None

    pred = preds[0]
    shape = orig_img.shape
    pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round()
    return pred.cpu().numpy()


def trim_holes(holes):
    """
    去掉尺寸异常的hole，大于平均值的1.2倍，或小于0.85倍的，忽略掉
    """
    if holes is None or len(holes) == 0:
        return []
    avg_w = 0.
    for _, _, xyxy in holes:
        avg_w += (xyxy[2] - xyxy[0]) / len(holes)

    arr = []
    for cls, conf, xyxy in holes:
        w = xyxy[2] - xyxy[0]
        if w > 1.2 * avg_w:
            continue
        if w < 0.85 * avg_w:
            continue

        arr.append((cls, conf, xyxy))

    return arr


def plot_one_box(x, img, color=color, label=None, line_thickness=1):
    # Plots one bounding box on image img
    # tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1  # line/font thickness
    tl = line_thickness
    # color = color or [random.randint(0, 255) for _ in range(3)]
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
    if label:
        tf = max(tl - 1, 1)  # font thickness
        t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
        c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
        # cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled

        # 标签再左上部
        # cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
        if label:
            cv2.putText(img, label, (c1[0] + 2, c1[1] + 14), 0, tl / 2, color, thickness=tf, lineType=cv2.LINE_AA)

        # 标签再中间
        # x1 = int((c1[0]+c2[0])/2)
        # y1 = int((c1[1]+c2[1])/2)
        # cv2.putText(img, label, (100, 100), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
        # cv2.putText(img, label, (40, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)


def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True,
              stride=stride):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)

    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)

    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

    dw /= 2  # divide padding into 2 sides
    dh /= 2

    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, ratio, (dw, dh)
